from data_loader.siamese_data_loader import SiameseDataLoader
from models.siamese_model import BiLSTMSiameseModel, SiameseCategoricalModel, SiameseMixedCNNBiLSTMDistanceModel,\
    SiameseMixedCategoricalModel, SiameseMixedBiLSTMDistanceModel, BiLSTMAttSiameseModel
from configs.siamese_config import Config
from utils.utils import make_submission

import os


def bilstm_siamese(level, test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'bilstm_siamese_' + level
    config.embedding_path = 'custom_data/%s_level' % level
    if level == 'word':
        config.max_len = config.max_len_word
    else:
        config.max_len = config.max_len_char

    print('Create the data generator.')
    data_loader = SiameseDataLoader(config,
                                    test_file=test_file_path,
                                    level=level)

    print('Create the model.')
    bilstm_siamese_model = BiLSTMSiameseModel(config, data_loader.get_vocabulary())

    print('Start predicting results.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    bilstm_siamese_model.load(best_weights)
    test_x_a, test_x_b = data_loader.get_test_data()
    test_results = bilstm_siamese_model.model.predict([test_x_a, test_x_b])
    test_y_pred = [test_result[0] < 0.65 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)


def bilstm_siamese_mixed(test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'bilstm_siamese_mixed'
    config.embedding_path_word = 'custom_data/word_level'
    config.embedding_path_char = 'custom_data/char_level'

    print('Create the data generator.')
    data_loader_word = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='word')

    data_loader_char = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='char')

    print('Create the model.')
    bilstm_siamese_mixed_model = SiameseMixedBiLSTMDistanceModel(config,
                                                                 data_loader_word.get_vocabulary(),
                                                                 data_loader_char.get_vocabulary())

    print('Start predicting results.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    bilstm_siamese_mixed_model.load(best_weights)
    test_x_a_word, test_x_b_word = data_loader_word.get_test_data()
    test_x_a_char, test_x_b_char = data_loader_char.get_test_data()
    test_results = bilstm_siamese_mixed_model.model.predict([test_x_a_word, test_x_b_word,
                                                             test_x_a_char, test_x_b_char])
    test_y_pred = [test_result[0] < 0.65 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)


def cnn_bilstm_siamese_mixed(test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'cnn_bilstm_siamese_mixed'
    config.embedding_path_word = 'custom_data/word_level'
    config.embedding_path_char = 'custom_data/char_level'

    print('Create the data generator.')
    data_loader_word = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='word')

    data_loader_char = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='char')

    print('Create the model.')
    cnn_bilstm_siamese_mixed_model = SiameseMixedCNNBiLSTMDistanceModel(config,
                                                                        data_loader_word.get_vocabulary(),
                                                                        data_loader_char.get_vocabulary())

    print('Start predicting results.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    cnn_bilstm_siamese_mixed_model.load(best_weights)
    test_x_a_word, test_x_b_word = data_loader_word.get_test_data()
    test_x_a_char, test_x_b_char = data_loader_char.get_test_data()
    test_results = cnn_bilstm_siamese_mixed_model.model.predict([test_x_a_word, test_x_b_word,
                                                                 test_x_a_char, test_x_b_char])
    test_y_pred = [test_result[0] < 0.65 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)


def siamese_categorical(level, test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'siamese_categorical_' + level
    config.embedding_path = 'custom_data/%s_level' % level
    if level == 'word':
        config.max_len = config.max_len_word
    else:
        config.max_len = config.max_len_char

    print('Create the data generator.')
    data_loader = SiameseDataLoader(config,
                                    test_file=test_file_path,
                                    level=level)

    print('Create the model.')
    siamese_categorical_model = SiameseCategoricalModel(config, data_loader.get_vocabulary())

    print('Start predicate results.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    siamese_categorical_model.load(best_weights)
    test_x_a, test_x_b = data_loader.get_test_data()
    test_results = siamese_categorical_model.model.predict([test_x_a, test_x_b])
    test_y_pred = [test_result[0] > 0.4 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)


def siamese_categorical_mixed(test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'siamese_categorical_mixed'
    config.embedding_path_word = 'custom_data/word_level'
    config.embedding_path_char = 'custom_data/char_level'

    print('Create the data generator.')
    data_loader_word = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='word')

    data_loader_char = SiameseDataLoader(config,
                                         test_file=test_file_path,
                                         level='char')

    print('Create the model.')
    siamese_categorical_model = SiameseMixedCategoricalModel(config,
                                                             data_loader_word.get_vocabulary(),
                                                             data_loader_char.get_vocabulary())

    print('Start evaluate the model.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    siamese_categorical_model.load(best_weights)
    test_x_a_word, test_x_b_word = data_loader_word.get_test_data()
    test_x_a_char, test_x_b_char = data_loader_char.get_test_data()
    test_results = siamese_categorical_model.model.predict([test_x_a_word, test_x_b_word,
                                                            test_x_a_char, test_x_b_char])
    test_y_pred = [test_result[0] > 0.4 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)


def bilstm_att_siamese(level, test_file_path, output_file_path):
    config = Config()
    config.exp_name = 'bilstm_att_siamese_' + level
    config.embedding_path = 'custom_data/%s_level' % level
    if level == 'word':
        config.max_len = config.max_len_word
    else:
        config.max_len = config.max_len_char

    print('Create the data generator.')
    data_loader = SiameseDataLoader(config,
                                    test_file=test_file_path,
                                    level=level)

    print('Create the model.')
    bilstm_att_siamese_model = BiLSTMAttSiameseModel(config, data_loader.get_vocabulary())

    print('Start predicting results.')
    best_weights = os.path.join('experiments', config.exp_name, 'checkpoints', config.exp_name + '.hdf5')
    bilstm_att_siamese_model.load(best_weights)
    test_x_a, test_x_b = data_loader.get_test_data()
    test_results = bilstm_att_siamese_model.model.predict([test_x_a, test_x_b])
    test_y_pred = [test_result[0] < 0.65 for test_result in test_results]
    make_submission(test_y_pred, output_file_path)
